import my_bonito.config as config
import random


def generate_random_task_list(task_list, proportion, base=100):
    if len(task_list) != len(proportion):
        raise ValueError("The length of task_list and proportion must be the same.")

    task_list = [config.SHORTFORM_TO_FULL_TASK_TYPES[task] for task in task_list]

    proportion = [int(prop) for prop in proportion]
    total_tasks = sum(proportion)
    # 按base划分，因为取模会有损失，所以base越大越准确，最终生成约base长度的列表，是按比例分布的各任务打乱版
    task_counts = {task: prop * base // total_tasks for task, prop in zip(task_list, proportion)}
    task_list_in_proportion = sum(([task] * count for task, count in task_counts.items()), [])
    random.shuffle(task_list_in_proportion)
    print('task_list_in_proportion:', task_list_in_proportion)
    return task_list_in_proportion


def generate_prompt_from_task_list(task_list, total_num):
    prompts = []
    task_num_one_request = config.NUM_OF_TASK_ONE_REQUEST
    num_prompts = total_num // task_num_one_request + (1 if total_num % task_num_one_request else 0)
    for i in range(num_prompts):
        prompt_start = i * task_num_one_request % len(task_list)
        prompt_end = prompt_start + task_num_one_request
        if prompt_end > len(task_list):
            tasks_slice = task_list[prompt_start:] + task_list[:prompt_end % len(task_list)]
        else:
            tasks_slice = task_list[prompt_start:prompt_end]

        task_str = "、".join(f"{task}" for task in tasks_slice)
        prompt = config.PROMPT_START+task_str+config.PROMPT_END
        prompts.append(prompt)
    return prompts


def generate_prompts(task_list, proportion, total_num, base=100):
    # base决定精度
    task_list_in_proportion = generate_random_task_list(task_list, proportion, base=base)
    prompts = generate_prompt_from_task_list(task_list_in_proportion, total_num)
    return prompts


if __name__ == "__main__":
    task_list_example = ["exqa", "ynqa", "mcqa", "nli", "qa"]
    proportion_example = ["4", "3", "3", "5", "12"]
    total_num = config.TOTAL_TASK_NUM
    prompts = generate_prompts(task_list_example, proportion_example, total_num)
    for prompt in prompts:
        print(prompt)